import json
import re
import ssl
from datetime import datetime

import requests
import websocket

refreshToken = ("eyJhbGciOiJIUzUxMiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJ1c2VyLWNlbnRlciIsImV4cCI6MTcxODI2MDg2NywiaWF0IjoxNzEw"
                "NDg0ODY3LCJqdGkiOiJjbnB1cjB1Y3A3ZjFodXI4dHY0ZyIsInR5cCI6InJlZnJlc2giLCJzdWIiOiJjbXIza3RlY3A3ZmR2cjFuZ"
                "jRoZyIsInNwYWNlX2lkIjoiY21yM2t0ZWNwN2ZkdnIxbmY0aDAiLCJhYnN0cmFjdF91c2VyX2lkIjoiY21yM2t0ZWNwN2ZkdnIxbmY"
                "0Z2cifQ.yvOhQLaifS0hjHMfPCUxrNY4ywz3q3N5j2DEx3WpcyYEm7GBMsrRrcmFAaNekoarTSZFXn4IQjN8AAtFdT4lAg")

attachments = ["coc9mehhmfreoh7n9o1g", "cntfchp98onppvieq3a0"]  # 附件id 帮助中心.txt rockflow_doc1.docx


class KimiClient:
    def __init__(self):
        self.accessToken = None

    def getAccessToken(self):
        url = "https://kimi.moonshot.cn/api/auth/token/refresh"
        # 获取当前时间
        now = datetime.utcnow()
        # 格式化日期时间字符串
        formatted_date = now.strftime('%a, %d %b %Y %H:%M:%S GMT')
        payload = {}
        headers = {
            'Host': 'kimi.moonshot.cn',
            'authorization': "Bearer %s" % refreshToken,
            'user-agent': 'okhttp/4.9.2',
            'if-modified-since': formatted_date
        }
        response = requests.request("GET", url, headers=headers, data=payload).text
        data = json.loads(response)
        return data['access_token']

    def createNewChat(self):
        url = "https://kimi.moonshot.cn/api/chat"
        payload = json.dumps({
            "name": "未命名会话"
        })
        headers = {
            'Host': 'kimi.moonshot.cn',
            'r-timezone': 'GMT',
            'x-msh-platform': 'android',
            'authorization': self.accessToken,
            'content-type': 'application / json',
            'user-agent': 'okhttp/4.9.2'
        }
        response = requests.request("POST", url, headers=headers, data=payload).text
        data = json.loads(response)
        if "error_type" in data and data["error_type"] == "auth.token.invalid":
            return None, response
        return data["id"], response

    def connentWebsockt(self, chatId, prompt, needDocument=True, debug=False):
        received_message = ""
        message_received = False
        ws_url = "wss://kimi.moonshot.cn/api/chat/stream/ws/%s?platform=android&token=%s" % (chatId, self.accessToken)
        ws = websocket.create_connection(ws_url, sslopt={"cert_reqs": ssl.CERT_NONE})
        documents = attachments
        if not needDocument:
            documents = []

        sendBody = {"action": "consult",
                    "consult_body":
                        {
                            "chat_id": chatId,
                            "data": {
                                "use_search": True,
                                "refs": documents,
                                "messages": [{"role": "user", "content": prompt}]
                            }
                        }
                    }
        ws.send(json.dumps(sendBody))
        # 在循环中持续接收数据
        while not message_received:
            try:
                # 接收数据
                result = re.match(r'data: (\{.*\})', ws.recv()).group(1)
                data = json.loads(result)
                if debug:
                    print(data)
                if data["event"] == "cmpl":
                    received_message += data["text"]
                if data["event"] == "error":
                    received_message += "接客接累了，下一次再来找我😭"
                elif data["event"] == "stream_done":
                    message_received = True
            except websocket.WebSocketException as e:
                if debug:
                    print(data)
                break
            except Exception as e:
                if debug:
                    print(data)
                # 处理连接关闭或其他异常情况
                break
        ws.close()
        received_message = received_message.replace("根据您提供", "根据帮助中心提供")
        return received_message

    def filterAnswer(self, answer):
        answer = answer.replace("Kimi", "bobbie")
        answer = answer.replace("月之暗面", "XXX")
        answer = answer.replace("Moonshot", "")
        answer = answer.replace("**", "")
        return answer

    def createNewChatId(self, isFirst=True):
        newChatId, result = self.createNewChat()
        if newChatId is None:
            if isFirst:
                self.accessToken = self.getAccessToken()
                return self.createNewChatId(False)
            else:
                return newChatId, result
        return newChatId, result

    def doPrompt(self, prompt, debug=False, needDocument=True):
        if self.accessToken is None:
            self.accessToken = self.getAccessToken()
        newChatId, result = self.createNewChatId()
        if newChatId is None:
            return "I am very tired.If you still want to continue, ask me again." + result
        answer = self.connentWebsockt(newChatId, prompt, needDocument, debug)
        answer = self.filterAnswer(answer)
        return answer
